import express from "express";
import { randomUUID } from "node:crypto";
import { McpServer, ResourceTemplate } from "@modelcontextprotocol/sdk/server/mcp.js";
import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/streamableHttp.js";
import { isInitializeRequest } from "@modelcontextprotocol/sdk/types.js"
import { z } from "zod";


const app = express();
app.use(express.json());

// Map to store transports by session ID
const transports = {};

// Handle POST requests for client-to-server communication
app.post('/mcp', async (req, res) => {
    // Check for existing session ID
    const sessionId = req.headers['mcp-session-id'];
    let transport;

    if (sessionId && transports[sessionId]) {
        // Reuse existing transport
        transport = transports[sessionId];
    } else if (!sessionId && isInitializeRequest(req.body)) {
        // New initialization request
        transport = new StreamableHTTPServerTransport({
            sessionIdGenerator: () => randomUUID(),
            onsessioninitialized: (sessionId) => {
                // Store the transport by session ID
                transports[sessionId] = transport;
            }
        });

        // Clean up transport when closed
        transport.onclose = () => {
            if (transport.sessionId) {
                delete transports[transport.sessionId];
            }
        };
        const server = new McpServer({
            name: "example-server",
            version: "1.0.0"
        });

        // ... set up server resources, tools, and prompts ...
        server.resource(
            "echo",
            new ResourceTemplate("echo://{message}", { list: undefined }),
            async (uri, { message }) => ({
                contents: [{
                    uri: uri.href,
                    text: `Resource echo: ${message}`
                }]
            })
        );

        server.tool("student_query",
            { age: z.number(), email: z.string() },
            async ({ }) => {
                return {
                    content: [
                        {
                            type: "text",
                            text: `查询结果为json格式，数据如下：[{ "id": 11, "name": "hl", "age": 18, "email": "fedhong.com", "created_at": "2025-06-04T02:00:53.000Z", "updated_at": "2025-06-04T02:00:53.000Z" }, { "id": 10, "name": "王十二", "age": 33, "email": "wangshier@example.com", "created_at": "2025-06-03T14:58:52.000Z", "updated_at": "2025-06-03T14:58:52.000Z" }, { "id": 9, "name": "郑十一", "age": 26, "email": "zhengshiyi@example.com", "created_at": "2025-06-03T14:58:52.000Z", "updated_at": "2025-06-03T14:58:52.000Z" }, { "id": 8, "name": "吴十", "age": 29, "email": "wushi@example.com", "created_at": "2025-06-03T14:58:52.000Z", "updated_at": "2025-06-03T14:58:52.000Z" }, { "id": 7, "name": "周九", "age": 31, "email": "zhoujiu@example.com", "created_at": "2025-06-03T14:58:52.000Z", "updated_at": "2025-06-03T14:58:52.000Z" }, { "id": 6, "name": "孙八", "age": 27, "email": "sunba@example.com", "created_at": "2025-06-03T14:58:52.000Z", "updated_at": "2025-06-03T14:58:52.000Z" }, { "id": 5, "name": "钱七", "age": 35, "email": "qianqi@example.com", "created_at": "2025-06-03T14:58:52.000Z", "updated_at": "2025-06-03T14:58:52.000Z" }, { "id": 4, "name": "赵六", "age": 22, "email": "zhaoliu@example.com", "created_at": "2025-06-03T14:58:52.000Z", "updated_at": "2025-06-03T14:58:52.000Z" }, { "id": 3, "name": "王五", "age": 28, "email": "wangwu@example.com", "created_at": "2025-06-03T14:58:52.000Z", "updated_at": "2025-06-03T14:58:52.000Z" }, { "id": 2, "name": "李四", "age": 30, "email": "lisi@example.com", "created_at": "2025-06-03T14:58:52.000Z", "updated_at": "2025-06-03T14:58:52.000Z" }]`
                        }
                    ]
                }
            }
        );

        server.tool(
            "echo",
            { message: z.string() },
            async ({ message }) => ({
                content: [{ type: "text", text: `Tool echo: ${message}` }]
            })
        );

        server.prompt(
            "echo",
            { message: z.string() },
            ({ message }) => ({
                messages: [{
                    role: "user",
                    content: {
                        type: "text",
                        text: `Please process this message: ${message}`
                    }
                }]
            })
        );

        // Connect to the MCP server
        await server.connect(transport);
        await transport.handlePostMessage(req, res, req.body);
    } else {
        // Invalid request
        res.status(400).json({
            jsonrpc: '2.0',
            error: {
                code: -32000,
                message: 'Bad Request: No valid session ID provided',
            },
            id: null,
        });
        return;
    }

    // Handle the request
    await transport.handleRequest(req, res, req.body);
});

// Reusable handler for GET and DELETE requests
const handleSessionRequest = async (req, res) => {
    const sessionId = req.headers['mcp-session-id'];
    if (!sessionId || !transports[sessionId]) {
        res.status(400).send('Invalid or missing session ID');
        return;
    }

    const transport = transports[sessionId];
    await transport.handleRequest(req, res);
};

// Handle GET requests for server-to-client notifications via SSE
app.get('/mcp', handleSessionRequest);

// Handle DELETE requests for session termination
app.delete('/mcp', handleSessionRequest);

// Start the server
const PORT = 3000;
app.listen(PORT, () => {
    console.log(`MCP Stateful Streamable HTTP Server listening on port ${PORT}`);
});